#!/usr/bin/env python3
import os
import re
from Bio import Entrez
from medline import parse

Entrez.email = "liqiming1914658215@gmail.com"                                      
Entrez.api_key = "c80ce212c7179f0bbfbd88495a91dd356708"

chineseLast = set(line.rstrip() for line in open('ChineseFamily.csv'))
chineseFirst = set(line.rstrip() for line in open('ChineseFirst.csv'))

journal_impact_dict = {}
for line in open("PubMed_Journal_Impact_factor.tsv"):
    toks = line.rstrip().split("\t")
    journal_impact_dict[toks[1]] = toks[2]


def isChinese(author):
    first = author["firstname"]
    last = author["lastname"]

    first = first.lower().replace("-", "")
    last = last.lower()
    if last in chineseLast and first in chineseFirst:
        return True
    else:
        return False

def Chines_authors(au):

    authors = {}

    first_au = au[0]
    last_au = au[-1]
    
    if isChinese(first_au) and  'China' in first_au['affiliation']:
        authors['first'] = first_au
    if isChinese(last_au) and  'China' in last_au['affiliation']:
        authors['last'] = last_au

    return authors

class QueryPubmed:
    def __init__(self, mesh, keywords, min_date, max_date, retmax=10000):
        self.mesh = mesh
        self.keywords = keywords
        self.min_date = min_date
        self.max_date = max_date
        self.retmax = retmax

    @property
    def query(self):
        return f"{self.mesh}[MeSH Major Topic] AND ({self.keywords}) AND {self.min_date}: {self.max_date}[dp]"

    def search(self):
        handle = Entrez.esearch(db="pubmed", term=self.query, retmax=self.retmax)
        record = Entrez.read(handle)
        return record["IdList"]

    def detail(self):
        idlist = self.search()
        count = len(idlist)
        print("Search '{}', get {} records.".format(self.query, count))
        handle = Entrez.efetch(db="pubmed", id=idlist, rettype="medline", retmode="text")
        records = parse(handle, journal_impact_dict, count)

        for record in records:
            impact_factor = journal_impact_dict[record.journal]
            authors = Chines_authors(record.authors)
            if not authors:
                continue

            article = [record.pubmed_id, record.title, record.journal, record.publication_date, record.abstract, record.keywords, impact_factor]
            yield article, record.authors


def main():
    results = QueryPubmed("Biology", "Bioinformatics OR Method OR Technology OR Genomics OR Application OR Database", '2020/01/01', '2020/12/31')
    count = 0
    with open("Bioinfomatics.txt", "w") as f:
        for article, authors in results.detail():
            count += 1
            print(f"{article[0]} saved.")
            f.write(f"Pubmed id: {article[0]}\n")
            f.write(f"Title: {article[1]}\n")
            f.write(f"Journal: {article[2]}\t{article[6]}\n")
            f.write(f"Publication date: {article[3]}\n")
            f.write("Authors:\n")
            for author in authors:
                if author['affiliation'] == 'NA':
                    f.write(f"    {author['fullname']}\n")
                else:
                    f.write(f"    {author['fullname']}  {author['affiliation']}\n")
            abstract = '\n    '.join(article[4])
            f.write(f"Abstract:\n    {abstract}\n")

            if article[5]:
                f.write(f"Keywords: {article[5]}\n")
            f.write("\n")
    print(f"Get {count} records.")


if __name__ == "__main__":
    main()